#include <iostream>
#include <vector>
#include <ctime>
#include <numeric>

#include "tensor.h"
#include "covariance.h"
#include "mtrand.h"
#include "dataset.h"
#include "pdf.h"
#include <vector>
#include <sstream>

#include "sdataset.h"
#include "spdf.h"

#include "TFile.h"
#include "TChain.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TProfile2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TF3.h"
#include "TLorentzVector.h"
#include "TGraphErrors.h"
#include "TGraph2DErrors.h"
#include "make_list.h"

using namespace std;
typedef vector<double> V1;
typedef valarray<Double_t> valD;
typedef valarray<double> valF;
typedef vector<valD> valD_List;

//initialize random number generator
MTRand_int32 irand(time(0));

int option=1;

Double_t func2D(Double_t x, Double_t y, int c)
{
  Double_t norm=1.0;
  Double_t epsilon=1.0;
  //normalization
  switch(c)
    {
    case 1:
      epsilon=0.001;
      norm= 1.0/0.139586;
      break;

    case 2:
      epsilon=0.002;
      norm= 1.0/0.138199;
      break;
    
    case 3:
      epsilon=0.005;
      norm= 1.0/0.134685;
      break;
    
    case 4:
      epsilon=0.01;
      norm= 1.0/0.130166;
      break;
    
    case 5:
      epsilon=0.02;
      norm= 1.0/0.123705;
      break;
      
    case 6:
      epsilon=0.5;
      norm= 1.0/0.0768968;
      break;

 
    default:
      throw;
    }

  Double_t xx=x-0.5;
  xx*=xx;
  
  Double_t yy=y-0.5;
  yy*=yy;
  
  Double_t xy=x-y;
  xy*=xy;
  xy*=xy;

  return norm* exp( -22.2222222*(xx+yy) - 1975.3*epsilon*xy);
}

//define a 2D function
Double_t testfunc_2D(Double_t* x, Double_t* par)
{
  return func2D(x[0], x[1], option);
}

pair<valD_List, valD>
GenerateMC( Double_t (*pfunc)(Double_t*, Double_t*), int nvar, int nevt=1000,
	    Double_t* factor=NULL)
{
  valD_List result(nevt,valarray<Double_t>(nvar));
  valD weight(nevt);

  for(int i=0; i<nevt; ++i)
    {
      //generate nvar random numbers
      for(int j=0; j<nvar; j++)
	{
	  if(factor==NULL)
	    result[i][j]=irand.rand_double();
	  else
	    result[i][j]=factor[j]*irand.rand_double();

	}
      //grab the weight
      weight[i]=pfunc(&result[i][0], NULL);  
    }
  //rescale the weights
  Double_t sumweight=weight.sum();
  
  sumweight=nevt/sumweight;

  for(int i=0; i<nevt; ++i)
    {
      weight[i]*=sumweight;
    }

  return make_pair(result, weight);
}

//function to create a graph
TGraph2D* MakeGraph2D(string str, valD& x, valD& y,
		      valD& value)
{
  if(x.size()!=value.size() || x.size()!=y.size())
    {
      cout<<"valarray sizes don't match in MakeGraph!!"<<endl;
      throw;	
    }
  
  TGraph2D* result;
  
  result=new TGraph2D(x.size(), &x[0], &y[0], &value[0]);
  
  result->SetTitle(str.c_str());
  result->SetName(str.c_str());
  result->SetLineWidth(2);
  
  return result;
}


int main(int argc, char **argv)
{

  if(argc < 2+1)
    {
      cout<<"Usage: ./example_qcd output.root [option]"<<endl;
      throw;
    }
  
  
  //open output file
  TFile outfile(argv[1], "RECREATE");

  //get the option
  option=atoi(argv[2]);

  const double PI=3.14159265359;

  //*********************
  // 2D testing
  // declare a root function
  
  cout<<"Staring 2D testing"<<endl;

  TF2 root_func2D("func2D", testfunc_2D, 0, 2.0, 0, 1.0, 2);
  
  //get the real function
  double num_2D=30.;
  valD x_2D(num_2D*num_2D);
  valD y_2D(num_2D*num_2D);
  valD func_2D(num_2D*num_2D);
  for(int i=0; i<num_2D; i++)
  for(int j=0; j<num_2D; j++)
    {
      int index=i*num_2D+j;
      x_2D[index]=i/num_2D;
      y_2D[index]=j/num_2D;
      func_2D[index]=root_func2D.Eval(x_2D[index], y_2D[index]);
    }

  TGraph2D* g2D_real=MakeGraph2D("real_2D", x_2D, y_2D, func_2D);
  g2D_real->Write();

  //get some MC events
  int nMC=2000000;
  pair<valD_List, valD> MC_func2D =
    GenerateMC(testfunc_2D, 2, nMC);
  
  valD_List value= MC_func2D.first;
  valD weights= MC_func2D.second;

  //initialize a DataSet
  DataSet data2D(make_list<double>(0.0, 0.0),
		 make_list<double>(1.0, 1.0),
		 make_list<double>(0.01, 0.01));

  //initialize 1D dataset
  DataSet datax(make_list<double>(0.0),
		make_list<double>(1.0),
		make_list<double>(0.01));

  DataSet datay(make_list<double>(0.0),
		make_list<double>(1.0),
		make_list<double>(0.01));

  
  //populate the DataSet with MC data
  //need to convert Double_t to double
  //also get correlation
  double mean_x=0;
  double mean_y=0;
  double mean_xy=0;
  double sum_weight=0;
  double mean_xx=0;
  double mean_yy=0;


  for(int j=0; j<value.size(); ++j)
    {
      data2D.Fill(value[j],
		  weights[j]);
      sum_weight+=weights[j];
      mean_x+=value[j][0]*weights[j];
      mean_y+=value[j][1]*weights[j];
      mean_xy+=value[j][0]*value[j][1]*weights[j];
      mean_xx+=value[j][0]*value[j][0]*weights[j];
      mean_yy+=value[j][1]*value[j][1]*weights[j];

      datax.Fill(make_list<double>(value[j][0]), weights[j]);
      datay.Fill(make_list<double>(value[j][1]), weights[j]);

    }

  sum_weight=1.0/sum_weight;
  mean_x*=sum_weight;
  mean_y*=sum_weight;
  mean_xy*=sum_weight;
  mean_xx*=sum_weight;
  mean_yy*=sum_weight;
  
  double correlation=mean_xy-mean_x*mean_y;
  cout<<"correlation is: "<<correlation/
    sqrt((mean_xx-mean_x*mean_x)*(mean_yy-mean_y*mean_y))
      <<endl;

  //now grab the PDF
  PDF pdf2D_p = data2D.ComputePDF();

  //read to file
  pdf2D_p.Write("PDFout.txt");

  PDF pdf2D;
  pdf2D.Read("PDFout.txt");

  pdf2D.Write("PDFout_new.txt");


  PDF pdfx = datax.ComputePDF();
  PDF pdfy = datay.ComputePDF();

  //now get the resulting graph
  valD test_2D(x_2D.size());
  for(int i=0; i<x_2D.size(); ++i)
    {
      double err;
      test_2D[i]=pdf2D(make_list<double>(x_2D[i],y_2D[i]), &err);
    }

  //get the graph
  TGraph2D* g2D_test=MakeGraph2D("test_2D", x_2D, y_2D, test_2D);

  //write the files
  g2D_test->Write();


  //now we will generate MC events from the 1D pdf
  //first declare some MC Generators
  PDF::MC_Generator MC_x = pdfx.Generator
    (make_list<PDF::MC_Flag>(PDF::OUTPUT));

  PDF::MC_Generator MC_y = pdfy.Generator
    (make_list<PDF::MC_Flag>(PDF::OUTPUT));

  vector<PDF::MC_Event>  MC_evt_x=
    MC_x.GenerateGrad(make_list<double>(), nMC);

  vector<PDF::MC_Event>  MC_evt_y=
    MC_y.GenerateGrad(make_list<double>(), nMC);

  //now reweight the events
  double rescale_reweight=0;
  for(unsigned int j=0; j<MC_evt_x.size(); ++j)
    {
      rescale_reweight +=
	(nMC*MC_evt_x[j].weight*nMC* MC_evt_y[j].weight
	 + correlation*
	 MC_evt_x[j].GetExtra()[0] * MC_evt_y[j].GetExtra()[0]);
    }

  rescale_reweight=1.0/rescale_reweight;

  DataSet data2D_corr(make_list<double>(0.0, 0.0),
		      make_list<double>(1.0, 1.0),
		      make_list<double>(0.01, 0.01));

  DataSet data2D_control(make_list<double>(0.0, 0.0),
			 make_list<double>(1.0, 1.0),
			 make_list<double>(0.01, 0.01));

  //redo kernel smoothing
  for(int j=0; j<MC_evt_x.size(); ++j)
    {
      double extra_factor=(nMC* MC_evt_x[j].weight *nMC* MC_evt_y[j].weight+ 
			   correlation*
			   MC_evt_x[j].GetExtra()[0] * 
			   MC_evt_y[j].GetExtra()[0]);
      
      /*
      cout<<"things: "<<nMC*MC_evt_x[j].weight<<","<<nMC*MC_evt_y[j].weight<<endl;
      cout<<"old: "<<nMC* MC_evt_x[j].weight *nMC* MC_evt_y[j].weight<<endl;
      cout<<"corr: "<<0.5*correlation*
	MC_evt_x[j].GetExtra()[0] * 
	MC_evt_y[j].GetExtra()[0]<<endl;
      */

      //cout<<"extra factor: "<<MC_evt_x[j].GetExtra()[0]<<","<<MC_evt_y[j].GetExtra()[0]<<endl;

      data2D_corr.Fill(make_list<double>(MC_evt_x[j].value[0], 
					 MC_evt_y[j].value[0]),
		       extra_factor);

      data2D_control.Fill(make_list<double>(MC_evt_x[j].value[0], 
					    MC_evt_y[j].value[0]),
			  MC_evt_x[j].weight * MC_evt_y[j].weight);
    }

  PDF pdf2D_corr = data2D_corr.ComputePDF();
  PDF pdf2D_control = data2D_control.ComputePDF();

  //also a plot to store deviations
  TH1D sig_2D_corr("sig_2D_corr", "Deviation for 2D test (corrected)", 50, -5, 5);

  TH1D sig_2D_control("sig_2D_control", 
		      "Deviation for 2D test (control)", 50, -5, 5);

  //get the values
  for(int i=0; i<x_2D.size(); ++i)
    {
      double err;
      test_2D[i]=pdf2D_corr(make_list<double>(x_2D[i],y_2D[i]), &err);

      //get the deviation
      if(func_2D[i]>0)
	sig_2D_corr.Fill((test_2D[i]-func_2D[i])/func_2D[i]);
    }

  //get the graph
  TGraph2D* g2D_test_corr=MakeGraph2D("test_corr_2D", x_2D, y_2D, test_2D);
  //write the files
  sig_2D_corr.Write();
  g2D_test_corr->Write();



  //get the values
  for(int i=0; i<x_2D.size(); ++i)
    {
      double err;
      test_2D[i]=pdf2D_control(make_list<double>(x_2D[i],y_2D[i]), &err);

      //get the deviation
      if(func_2D[i]>0)
	sig_2D_control.Fill((test_2D[i]-func_2D[i])/func_2D[i]);
    }

  //get the graph
  TGraph2D* g2D_test_control=MakeGraph2D("test_control_2D", x_2D, y_2D, test_2D);


  //write the files
  sig_2D_control.Write();
  g2D_test_control->Write();




  
  //declare some histograms
  map<string, TH1*> hist;
  /*
  hist["T21_MC_jetmass"]=new TH1D("T21_jetmass", 
				  "T21 from jetmass PDF", 20, 0, 1.0);
  hist["T21_MCerr_jetmass"]=new TH1D("T21err_jetmass", 
				  "T21 err from jetmass PDF", 20, 0, 1.0);
  */  
  
  //write all histograms
  map<string, TH1*>::iterator it;
  for(it=hist.begin(); it!=hist.end(); ++it)
    {
      //see if the name has MC but no MCerr
      if(it->first.find("MC")!=string::npos &&
	 it->first.find("MCerr")==string::npos)
	{
	  //try to see if the corresponding MCerr histogram exists
	  string str=it->first;
	  str.replace(str.find("MC"), 2, "MCerr");
	  TH1* ch=hist[str];
	  if(ch!=NULL)
	    {
	      //found a corresponding histogram
	      for(int i=1; i<=it->second->GetNbinsX(); ++i)
		{
		  it->second->SetBinError
		    (i, ch->GetBinContent(i));
		}

	    }
	}

      it->second->Write();
    }

  outfile.Close();
  cout<<"Program ends."<<endl;
  return 1;
}
  
